import numpy as np
from sklearn.metrics import classification_report
from sentence_transformers import SentenceTransformer, util
from preprocess import get_data, model_save_path, result_save_path, tolerance, epochs
import pandas as pd
import matplotlib.pyplot as plt


def f():

    model = SentenceTransformer(model_path)

    embedding1 = model.encode(s1, convert_to_tensor=True)
    embedding2 = model.encode(s2, convert_to_tensor=True)
    # ④预测准备工作：初始化所有句子标签为0
    pre_labels = [0] * len(s1)
    predict_file = open(result_save_path, 'w', encoding="utf-8")

    right = 0
    # label = np.zeros(10)
    for i in range(len(s1)):
        similarity = util.cos_sim(embedding1[i], embedding2[i])
        # if 0 <= similarity < 0.1:
        #     label[0] += 1
        # elif 0.1 <= similarity < 0.2:
        #     label[1] += 1
        # elif 0.2 <= similarity < 0.3:
        #     label[2] += 1
        # elif 0.3 <= similarity < 0.4:
        #     label[3] += 1
        # elif 0.4 <= similarity < 0.5:
        #     label[4] += 1
        # elif 0.5 <= similarity < 0.6:
        #     label[5] += 1
        # elif 0.6 <= similarity < 0.7:
        #     label[6] += 1
        # elif 0.7 <= similarity < 0.8:
        #     label[7] += 1
        # elif 0.8 <= similarity < 0.9:
        #     label[8] += 1
        # elif 0.9 <= similarity < 1:
        #     label[9] += 1
        pre_labels[i] = round(similarity.item(), 3)
        if abs(y_test[i] - pre_labels[i]) < tolerance:
            right += 1

        predict_file.write(s1[i] + ' ' +
                           s2[i] + ' ' +
                           str(y_test[i]) + ' ' +
                           str(pre_labels[i]) + '\n')

    print(model_path)
    # print(label / len(s1))
    acc = round(right / len(s1) * 100, 2)
    predict_file.write('准确率为' + str(round(right / len(s1) * 100, 2)) + '%')
    print('准确率为' + str(acc) + '%')
    predict_file.close()
    return acc


# 获取到评估结果后画图
def draw():
    x = []
    y1 = []
    y2 = []
    # 转化数据形式
    df = pd.DataFrame([x, y1, y2]).T
    # 对列重新命名
    df.columns = ['Cosine Similarity', 'short', 'long']
    # 数据写入图像，命名图例
    plt.title("Textual Similarity")
    plt.plot(df['Cosine Similarity'], df['short'], label='short')
    plt.plot(df['Cosine Similarity'], df['long'], label='long')
    plt.xlabel('Cosine Similarity')
    plt.ylabel('Probability')
    plt.legend()
    plt.savefig('./graph/' + 'Textual Similarity' + '.png', dpi=1600)


if __name__ == '__main__':

    # for i in range(1, epochs):
    #     model_path = './model_hub/resultModel_kaggle_short_checkpoints/' + str(i * 215)
    #     x_train, x_test, y_train, y_test = get_data()
    #     s1 = np.array(x_test)[:, 0]
    #     s2 = np.array(x_test)[:, 1]
    #     accuracy.append(f(x_train, x_test, y_train, y_test, s1, s2, model_path))

    model_path = './model_hub/resultModel_kaggle_long' + ''

    x_train, x_test, y_train, y_test = get_data('./data/kaggle_long.csv')
    s1 = np.array(x_test)[:, 0]
    s2 = np.array(x_test)[:, 1]
    f()

    # draw()
